{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b026af2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import models\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b960457",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\24566\\AppData\\Roaming\\Python\\Python312\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "C:\\Users\\24566\\AppData\\Roaming\\Python\\Python312\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet18_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet18_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 载入模型结构\n",
    "model = models.resnet18(pretrained=True) # 载入预训练模型\n",
    "model.fc = nn.Linear(model.fc.in_features, 81) # 81类\n",
    "model.fc\n",
    "model  = model.to(device)\n",
    "\n",
    "\n",
    "# 加载模型权重\n",
    "state_dict = torch.load(r'checkpoint\\best-0.869.pth') \n",
    "model.load_state_dict(state_dict)\n",
    "model.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "caa919ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 224, 224])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(1, 3, 224, 224).to(device)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "792f7e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \n",
    "with torch.no_grad(): # 关闭梯度\n",
    "    torch.onnx.export(\n",
    "        model,                   # 要转换的模型，必须是eval 模式\n",
    "        x,                       # 模型出入示例\n",
    "        'resnet18_fruit81.onnx', # 导出的 ONNX 文件名\n",
    "        opset_version=15,        # ONNX 算子集版本\n",
    "        input_names=['input'],   # 输入节点名称\n",
    "        output_names=['output']  # 输出节点名称\n",
    "    ) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecf34e3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
